package org.zjvis.datascience.service.dag;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import net.zjvis.lab.nebula.backend.lib.algorithm.conf.Configuration;
import net.zjvis.lab.nebula.backend.lib.algorithm.conf.algorithm.AlgorithmHelper;
import net.zjvis.lab.nebula.backend.lib.algorithm.conf.connector.Postgresql;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.WebClient;
import org.zjvis.datascience.common.dto.TaskInstanceDTO;
import org.zjvis.datascience.common.enums.AlgPyEnum;
import org.zjvis.datascience.common.enums.TaskTypeEnum;
import org.zjvis.datascience.common.model.Column;
import org.zjvis.datascience.common.util.RestTemplateUtil;
import org.zjvis.datascience.common.util.ToolUtil;
import org.zjvis.datascience.service.TaskInstanceService;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.stream.Collectors;

/**
 * @description Flask任务调度器
 * @date 2021-12-24
 */
public class FlaskSubmitRunner implements Callable<TaskRunnerResult> {

    private final static Logger logger = LoggerFactory.getLogger(FlaskSubmitRunner.class);

    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    private String errorTpl = "{\"status\":500, \"error_msg\":\"%s\"}";
    private String emptyTpl = "{\"status\":0, \"error_msg\":\"%s\"}";

    private TaskInstanceDTO instance;

    private RestTemplateUtil restTemplateUtil;

    private TaskInstanceService taskInstanceService;

    public FlaskSubmitRunner(RestTemplateUtil restTemplateUtil,
                             TaskInstanceService taskInstanceService, TaskInstanceDTO instance) {
        this.restTemplateUtil = restTemplateUtil;
        this.instance = instance;
        this.taskInstanceService = taskInstanceService;
    }

    @Override
    public TaskRunnerResult call() throws Exception {
        TaskRunnerResult result = null;
        if (instance.hasPrecautionaryError()) {
            return new TaskRunnerResult(500, String.format(errorTpl, "error happens when init stage."));
        }
        try {
            if (instance.getType().equals(TaskTypeEnum.TASK_TYPE_ALGOPY.getVal())) {
                this.syncInAndOutput();
                if (AlgPyEnum.getEnumByVal(JSONObject.parseObject(instance.getDataJson()).getJSONObject("inputInfo").getInteger("algType")).isNeedAlgo()) {
                    result = this.invoke();
                } else {
                    result = this.submit();
                }
            } else if (instance.getType().equals(TaskTypeEnum.TASK_TYPE_MODEL.getVal())) {
                result = this.exec();
            }

        } catch (Exception e) {
            result = new TaskRunnerResult(500,
                    String.format(errorTpl, e.getMessage().replaceAll("\"", "'")));
        }
        return result;
    }

    private void syncInAndOutput() {
        JSONObject jsonObject = JSONObject.parseObject(this.instance.getDataJson());
        JSONObject inputInfo = jsonObject.getJSONObject("inputInfo");
        List<String> inputCols = inputInfo.getJSONArray("input").getJSONObject(0)
                .getJSONArray("tableCols").toJavaList(String.class);
        inputInfo.getJSONArray("output").getJSONObject(0).put("tableCols", inputCols);

        List<String> columnTypes = inputInfo.getJSONArray("input").getJSONObject(0)
                .getJSONArray("columnTypes").toJavaList(String.class);
        inputInfo.getJSONArray("output").getJSONObject(0).put("columnTypes", columnTypes);
        jsonObject.put("inputInfo", inputInfo);
        instance.setDataJson(jsonObject.toJSONString());
    }

    /**
     * 针对有sql语句 或者 appArgs 命令的任务
     *
     * @return
     */
    private TaskRunnerResult exec() {
        return null;
    }

    public TaskRunnerResult invoke() throws JsonProcessingException, InstantiationException, IllegalAccessException {
        Long taskId = instance.getTaskId();
        logger.info("[FlaskSubmitRunner][invoke][{}][{}] start", taskId, instance.getId());
        JSONObject dataJson = JSONObject.parseObject(instance.getDataJson())
                .getJSONObject("inputInfo");
        JSONObject formData = dataJson.getJSONArray("setParams").getJSONObject(0)
                .getJSONObject("formData");

        logger.info("[FlaskSubmitRunner][invoke][{}][{}] dataJson={}", taskId, instance.getId(), dataJson);
        JSONObject input = dataJson.getJSONArray("input").getJSONObject(0);
        String source = input.getString("tableName");
        if (source.endsWith("_")) {
            Long parentTimeStamp = dataJson.getJSONArray("parentTimeStamps").getLong(0);
            source = source + parentTimeStamp;
        }
        int algoVal = dataJson.getInteger("algType");
        String algoEngName = AlgPyEnum.getEnumByVal(algoVal).getName();

        formData.put("taskId", taskId.toString());
        formData.put("source", source);
        formData.put("instanceId", instance.getId().toString());

        Postgresql inputDB = Postgresql.getDefault();
        inputDB.setReadSql("select * from " + ToolUtil.alignTableName(source, 0L));

        String outputTableName = String.format("%s_%d_%d", algoEngName, taskId, System.currentTimeMillis());
        String originalOutputTableName = "pipeline." + outputTableName;
        Postgresql outputDB = Postgresql.getDefault().toBuilder()
                .schema("pipeline")
                .ifWriteTableExists("replace")
                .writeTableName(outputTableName)
                .build();

        Configuration configuration = Configuration.builder()
                .connectorConf(Configuration.ConnectorConf.builder()
                        .input(inputDB)
                        .output(outputDB)
                        .build())
                .algorithm(AlgorithmHelper.parseFromJsonObject(dataJson))
            .taskId(taskId)
            .algorithmName(AlgorithmHelper.algorithmNameMap.get(algoEngName))
                .build();
        logger.info("[FlaskSubmitRunner][invoke][{}][{}] configuration={} ", taskId, instance.getId(), JSONObject.toJSONString(configuration));
        String string = OBJECT_MAPPER.writeValueAsString(configuration);
        String s = restTemplateUtil.invokeAlgo(string);

        boolean error = false;
        if (s.toLowerCase().contains("exception")) {
            instance.setLogInfo(String.format(errorTpl, s));
            error = true;
        } else {
            instance.setLogInfo(s);
        }

        logger.info("[FlaskSubmitRunner][invoke][{}][{}] call algo service finished, result={}",
                taskId, instance.getId(), s);
        try {
            instance = updateDataJson(instance, originalOutputTableName);
        } catch (Exception e) {
            logger.error("[FlaskSubmitRunner][invoke][{}][{}] update dataJson failed, error={}",
                    taskId, instance.getId(), JSON.toJSONString(e));
        }


        instance.setProgress(100);
        taskInstanceService.update(instance);

        logger.info("[FlaskSubmitRunner][invoke][{}][{}] finished, instance={} ", taskId, instance.getId(), instance.toString());
        return error ? TaskRunnerResult.fail(s) : TaskRunnerResult.ok(s);
    }

    private TaskInstanceDTO updateDataJson(TaskInstanceDTO taskInstance, String tableName) {
        logger.info("[FlaskSubmitRunner][invoke][updateDataJson][{}][{}] start, tableName={}, original taskInstance={} ",
                taskInstance.getTaskId(), taskInstance.getId(), tableName, JSON.toJSONString(taskInstance));
        String sql = String.format("SELECT * FROM %s", tableName);
        List<Column> columnList = taskInstanceService.queryDataColumn("", sql);
        String parentTaskInstanceId = taskInstance.getParentId();
        TaskInstanceDTO parentTaskInstance = taskInstanceService.queryById(Long.parseLong(parentTaskInstanceId));
        if (parentTaskInstance == null) {
            logger.error("[FlaskSubmitRunner][invoke][{}][{}] cannot find parent taskInstance by Id={}, current instance Id={}",
                    taskInstance.getTaskId(), taskInstance.getId(), parentTaskInstanceId, taskInstance.getId());
            return taskInstance;
        }
        logger.info("[FlaskSubmitRunner][invoke][updateDataJson][{}][{}]  columnList={}",
                taskInstance.getTaskId(), taskInstance.getId(), JSON.toJSONString(columnList));
        JSONObject originalDataJSON = JSONObject.parseObject(taskInstance.getDataJson());
        JSONObject dataJSON = originalDataJSON.getJSONObject("inputInfo");
        JSONObject parentDataJSON = JSONObject.parseObject(parentTaskInstance.getDataJson()).getJSONObject("inputInfo");
        logger.info("[FlaskSubmitRunner][invoke][updateDataJson][{}][{}]  original dataJson={}, parentDataJson={}",
                taskInstance.getTaskId(), taskInstance.getId(), JSON.toJSONString(dataJSON), JSON.toJSONString(parentDataJSON));
        JSONArray dataJSONOutput = parentDataJSON.getJSONArray("output");
        if (dataJSONOutput == null || dataJSONOutput.size() == 0) {
            return taskInstance;
        }
        List<String> newColumns = new ArrayList();
        List<String> columnNameList = new ArrayList<>();
        List<String> oldColumns = dataJSONOutput.getJSONObject(0).getJSONArray("tableCols").toJavaList(String.class);
        if (columnList != null) {
            newColumns = columnList.stream().filter(newColumn -> !oldColumns.contains(newColumn.getName())).map(Column::getName).collect(Collectors.toList());
            columnNameList = columnList.stream().map(Column::getName).collect(Collectors.toList());
            logger.info("[FlaskSubmitRunner][invoke][updateDataJson][{}][{}]  newColumns={}, columnNameList={}",
                    taskInstance.getTaskId(), taskInstance.getId(), JSON.toJSONString(newColumns), JSON.toJSONString(columnNameList));
        }
        dataJSONOutput.getJSONObject(0).put("tableCols", JSON.parseArray(JSON.toJSONString(columnNameList)));

        JSONObject newSemantic = new JSONObject();
        JSONObject oldSemantic = dataJSONOutput.getJSONObject(0).getJSONObject("semantic");
        Iterator iterator = oldSemantic.keySet().iterator();
        while (iterator.hasNext()) {
            String key = (String) iterator.next();
            String value = oldSemantic.getString(key);
            String newKey = key.replace(" ", "_").replace("/", "_").replace("-", "_");
            newSemantic.put(newKey, value);
        }
        dataJSONOutput.getJSONObject(0).put("semantic", newSemantic);
        dataJSONOutput.getJSONObject(0).put("tableName", tableName);

        for (String newColumn : newColumns) {
            dataJSONOutput.getJSONObject(0).getJSONObject("semantic").put(newColumn, "null");
            int index = dataJSONOutput.getJSONObject(0).getJSONArray("tableCols").indexOf(newColumn);
            dataJSONOutput.getJSONObject(0).getJSONArray("columnTypes").add(index, "REAL");
        }
        logger.info("[FlaskSubmitRunner][invoke][updateDataJson][{}][{}]  newSemantic={}, dataJSONOutput={}",
                taskInstance.getTaskId(), taskInstance.getId(), JSON.toJSONString(newSemantic), JSON.toJSONString(dataJSONOutput));

        dataJSON.put("output", dataJSONOutput);
        originalDataJSON.put("inputInfo", dataJSON);
        taskInstance.setDataJson(JSONObject.toJSONString(originalDataJSON));
        logger.info("[FlaskSubmitRunner][invoke][updateDataJson][{}][{}]  final dataJSON={}, updated instance={}",
                taskInstance.getTaskId(), taskInstance.getId(), JSON.toJSONString(originalDataJSON), JSON.toJSONString(taskInstance));
        return taskInstance;
    }

    /**
     * @return
     */
    public TaskRunnerResult submit() throws JsonProcessingException {
        Long taskId = instance.getTaskId();
        logger.warn("---------> {} ", taskId);
        JSONObject dataJson = JSONObject.parseObject(instance.getDataJson())
                .getJSONObject("inputInfo");
        JSONObject formData = dataJson.getJSONArray("setParams").getJSONObject(0)
                .getJSONObject("formData");

        logger.warn("---------> dataJson: {} ", dataJson);
        JSONObject input = dataJson.getJSONArray("input").getJSONObject(0);
        String source = input.getString("tableName");
        if (source.endsWith("_")) {
            Long parentTimeStamp = dataJson.getJSONArray("parentTimeStamps").getLong(0);
            source = source + parentTimeStamp;
        }
        int algoVal = dataJson.getInteger("algType");
        String algoEngName = AlgPyEnum.getEnumByVal(algoVal).getName();

        formData.put("taskId", taskId.toString());
        formData.put("target",
                String.format("pipeline.%s_%d_%d", algoEngName, taskId, System.currentTimeMillis()));
        formData.put("source", source);
        formData.put("instanceId", instance.getId().toString());

        String flaskServer = restTemplateUtil.getFlaskServer();
        String url = String.format("%s/%s", flaskServer, algoEngName);
        logger.warn("---------> url: {} ", url);
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_FORM_URLENCODED);

        MultiValueMap<String, String> formDataMap = new LinkedMultiValueMap<String, String>();
        for (String keyStr : formData.keySet()) {
            formDataMap.add(keyStr, formData.getString(keyStr));
        }
        logger.warn("formData -> {}", formDataMap.toString());
        WebClient webClient = WebClient.create();
        Flux<String> flux = webClient.post().uri(url)
                .body(BodyInserters.fromFormData(formDataMap)).retrieve().bodyToFlux(String.class);

        JSONObject resJson = JSONObject.parseObject(flux.blockFirst());
        boolean error = false;
        if (resJson.containsKey("error_msg")) {
            instance.setLogInfo(String.format(errorTpl, resJson.getString("error_msg")));
            error = true;
        } else {
            instance.setLogInfo(resJson.toJSONString());
        }
        instance.setProgress(100);
        taskInstanceService.update(instance);
        logger.warn("---------> instance: {} ", instance.toString());
        return error ? TaskRunnerResult.fail(resJson.getString("error_msg"))
                : TaskRunnerResult.ok(resJson.toJSONString());
    }

}
